package com.example.llm;

import javax.swing.*;
import java.awt.*;
import java.awt.event.*;
import java.io.*;
import java.util.prefs.*;
import javax.swing.text.BadLocationException;
import javax.swing.text.Document;
import java.util.List;
import java.util.ArrayList;
import reactor.core.Disposable;

import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;

public class MyllmClient extends JFrame {
    private JTextArea chatArea;
    private JTextField inputField;
    private JButton sendButton;
    private JButton settingsButton;
    private String modelName;
    private Preferences prefs;
    private static final String MODEL_PREF = "ollama_model";
    private boolean showDetails = false;
    private OllamaClient ollamaClient;
    private ChatSessionManager sessionManager;
    private JList<ChatSession> sessionList;
    private DefaultListModel<ChatSession> sessionListModel;
    
    // 配置项
    private String apiEndpoint = "http://127.0.0.1:11434";
    private String selectedModel = "deepseek:latest";
    private int maxContext = 20;
    private double temperature = 0.7;
    
    // 显示相关的配置
    private boolean showCharCount = false;
    private boolean showTokenCount = false;
    private boolean showTokenCost = true;
    private boolean showModelName = true;
    private boolean showTimestamp = false;
    private boolean showFirstTokenTime = false;
    private boolean autoFoldCodeBlock = true;

    private long startTime;
    private long firstTokenTime;
    private int totalTokens;
    private int totalChars;
    
    // 对话相关的配置
    private String systemPrompt = "You are a helpful assistant.";
    private boolean autoGenerateTitle = true;
    private boolean spellCheck = true;
    private boolean markdownRendering = true;
    private boolean latexRendering = true;
    private boolean mermaidRendering = true;
    private boolean injectDefaultData = true;
    private boolean autoGenerateArtifacts = false;
    private boolean longTextAsFile = true;

    private boolean isAIResponseStarted = false;
    private Disposable currentSubscription;
    private String currentSessionId;
    
    private StringBuilder messageBuffer = new StringBuilder();
    private static final int BUFFER_SIZE = 50; // 减小缓冲区大小，提高更新频率
    private Timer updateTimer;
    private volatile boolean isUpdating = false;
    private AtomicBoolean isFirstResponse = new AtomicBoolean(true);
    private final ExecutorService executorService = Executors.newSingleThreadExecutor();
    
    public MyllmClient() {
        super("myllm Chat");
        prefs = Preferences.userNodeForPackage(MyllmClient.class);
        sessionManager = new ChatSessionManager();
        loadSettings();
        initializeUI();
        ollamaClient = new OllamaClient(apiEndpoint);
    }
    
    private void initializeUI() {
        setDefaultCloseOperation(JFrame.EXIT_ON_CLOSE);
        setSize(1200, 800);
        setLocationRelativeTo(null);

        // 创建左侧会话列表面板
        sessionListModel = new DefaultListModel<>();
        sessionList = new JList<>(sessionListModel);
        sessionList.setCellRenderer(new SessionListCellRenderer());
        updateSessionList();
        
        // 添加右键菜单
        JPopupMenu popupMenu = new JPopupMenu();
        JMenuItem deleteItem = new JMenuItem("删除会话");
        deleteItem.addActionListener(e -> {
            ChatSession selected = sessionList.getSelectedValue();
            if (selected != null) {
                int confirm = JOptionPane.showConfirmDialog(
                    this,
                    "确定要删除会话 \"" + selected.getTitle() + "\" 吗？",
                    "删除会话",
                    JOptionPane.YES_NO_OPTION,
                    JOptionPane.WARNING_MESSAGE
                );
                if (confirm == JOptionPane.YES_OPTION) {
                    // 如果删除的是当前会话，先切换到其他会话
                    if (selected == sessionManager.getCurrentSession()) {
                        List<ChatSession> allSessions = sessionManager.getAllSessions();
                        for (ChatSession session : allSessions) {
                            if (session != selected) {
                                switchSession(session);
                                break;
                            }
                        }
                    }
                    
                    // 删除会话
                    sessionManager.deleteSession(selected);
                    
                    // 更新UI
                    updateSessionList();
                    chatArea.setText("");
                    
                    // 如果没有会话了，创建新会话
                    if (sessionManager.getAllSessions().isEmpty()) {
                        createNewSession();
                    }
                    
                    // 确保输入框状态正确
                    inputField.setEnabled(true);
                    sendButton.setEnabled(true);
                    inputField.requestFocus();
                }
            }
        });
        popupMenu.add(deleteItem);
        
        // 添加鼠标监听器
        sessionList.addMouseListener(new MouseAdapter() {
            private void showPopupMenu(MouseEvent e) {
                if (e.isPopupTrigger()) {
                    int index = sessionList.locationToIndex(e.getPoint());
                    if (index >= 0) {
                        sessionList.setSelectedIndex(index);
                        popupMenu.show(e.getComponent(), e.getX(), e.getY());
                    }
                }
            }
            
            @Override
            public void mousePressed(MouseEvent e) {
                showPopupMenu(e);
            }
            
            @Override
            public void mouseReleased(MouseEvent e) {
                showPopupMenu(e);
            }
            
            @Override
            public void mouseClicked(MouseEvent e) {
                if (SwingUtilities.isRightMouseButton(e)) {
                    int index = sessionList.locationToIndex(e.getPoint());
                    if (index >= 0) {
                        sessionList.setSelectedIndex(index);
                        popupMenu.show(e.getComponent(), e.getX(), e.getY());
                    }
                }
            }
        });
        
        JPanel leftPanel = new JPanel(new BorderLayout());
        JButton newSessionButton = new JButton("新建会话");
        newSessionButton.addActionListener(e -> createNewSession());
        leftPanel.add(newSessionButton, BorderLayout.NORTH);
        leftPanel.add(new JScrollPane(sessionList), BorderLayout.CENTER);
        
        sessionList.addListSelectionListener(e -> {
            if (!e.getValueIsAdjusting()) {
                ChatSession selected = sessionList.getSelectedValue();
                if (selected != null && !selected.getId().equals(currentSessionId)) {
                    switchSession(selected);
                }
            }
        });

        // 创建右侧聊天面板
        JPanel rightPanel = new JPanel(new BorderLayout());
        chatArea = new JTextArea();
        chatArea.setEditable(false);
        chatArea.setLineWrap(true);
        chatArea.setWrapStyleWord(true);
        
        JPanel inputPanel = new JPanel(new BorderLayout());
        inputField = new JTextField();
        sendButton = new JButton("发送");
        settingsButton = new JButton("设置");
        
        JPanel buttonPanel = new JPanel(new FlowLayout(FlowLayout.RIGHT));
        buttonPanel.add(settingsButton);
        
        // 添加性能监控按钮
        JButton perfButton = new JButton("性能监控");
        perfButton.addActionListener(e -> showPerformanceMonitor());
        buttonPanel.add(perfButton);
        
        buttonPanel.add(sendButton);
        
        inputPanel.add(inputField, BorderLayout.CENTER);
        inputPanel.add(buttonPanel, BorderLayout.EAST);
        
        rightPanel.add(new JScrollPane(chatArea), BorderLayout.CENTER);
        rightPanel.add(inputPanel, BorderLayout.SOUTH);

        // 创建分割面板
        JSplitPane splitPane = new JSplitPane(JSplitPane.HORIZONTAL_SPLIT, leftPanel, rightPanel);
        splitPane.setDividerLocation(250);
        add(splitPane);

        // 添加事件监听器
        inputField.addActionListener(e -> sendMessage());
        sendButton.addActionListener(e -> sendMessage());
        settingsButton.addActionListener(e -> showSettings());

        // 加载当前会话的消息
        loadCurrentSession();

        // 初始化定时器用于批量更新UI，减少更新间隔
        updateTimer = new Timer(50, e -> {
            if (!isUpdating && messageBuffer.length() > 0) {
                isUpdating = true;
                SwingUtilities.invokeLater(() -> {
                    String content = messageBuffer.toString();
                    messageBuffer.setLength(0);
                    if (isFirstResponse.get()) {
                        appendToChat("AI", content);
                        isFirstResponse.set(false);
                    } else {
                        appendMessage(content);
                    }
                    isUpdating = false;
                });
            }
        });
        updateTimer.start();
    }
    
    private void loadSettings() {
        apiEndpoint = prefs.get("apiEndpoint", "http://127.0.0.1:11434");
        selectedModel = prefs.get("selectedModel", "deepseek:latest");
        maxContext = prefs.getInt("maxContext", 20);
        temperature = prefs.getDouble("temperature", 0.7);
        
        // 加载显示相关的配置
        showCharCount = prefs.getBoolean("showCharCount", false);
        showTokenCount = prefs.getBoolean("showTokenCount", false);
        showTokenCost = prefs.getBoolean("showTokenCost", true);
        showModelName = prefs.getBoolean("showModelName", true);
        showTimestamp = prefs.getBoolean("showTimestamp", false);
        showFirstTokenTime = prefs.getBoolean("showFirstTokenTime", false);
        autoFoldCodeBlock = prefs.getBoolean("autoFoldCodeBlock", true);

        // 加载对话相关的配置
        systemPrompt = prefs.get("systemPrompt", "You are a helpful assistant.");
        autoGenerateTitle = prefs.getBoolean("autoGenerateTitle", true);
        spellCheck = prefs.getBoolean("spellCheck", true);
        markdownRendering = prefs.getBoolean("markdownRendering", true);
        latexRendering = prefs.getBoolean("latexRendering", true);
        mermaidRendering = prefs.getBoolean("mermaidRendering", true);
        injectDefaultData = prefs.getBoolean("injectDefaultData", true);
        autoGenerateArtifacts = prefs.getBoolean("autoGenerateArtifacts", false);
        longTextAsFile = prefs.getBoolean("longTextAsFile", true);
    }

    private void saveSettings() {
        prefs.put("apiEndpoint", apiEndpoint);
        prefs.put("selectedModel", selectedModel);
        prefs.putInt("maxContext", maxContext);
        prefs.putDouble("temperature", temperature);
        
        // 保存显示相关的配置
        prefs.putBoolean("showCharCount", showCharCount);
        prefs.putBoolean("showTokenCount", showTokenCount);
        prefs.putBoolean("showTokenCost", showTokenCost);
        prefs.putBoolean("showModelName", showModelName);
        prefs.putBoolean("showTimestamp", showTimestamp);
        prefs.putBoolean("showFirstTokenTime", showFirstTokenTime);
        prefs.putBoolean("autoFoldCodeBlock", autoFoldCodeBlock);

        // 保存对话相关的配置
        prefs.put("systemPrompt", systemPrompt);
        prefs.putBoolean("autoGenerateTitle", autoGenerateTitle);
        prefs.putBoolean("spellCheck", spellCheck);
        prefs.putBoolean("markdownRendering", markdownRendering);
        prefs.putBoolean("latexRendering", latexRendering);
        prefs.putBoolean("mermaidRendering", mermaidRendering);
        prefs.putBoolean("injectDefaultData", injectDefaultData);
        prefs.putBoolean("autoGenerateArtifacts", autoGenerateArtifacts);
        prefs.putBoolean("longTextAsFile", longTextAsFile);
    }
    
    private void appendDebugInfo(String info) {
        if (showDetails) {
            SwingUtilities.invokeLater(() -> {
                chatArea.append("[调试信息] " + info + "\n");
                chatArea.setCaretPosition(chatArea.getDocument().getLength());
            });
        }
    }

    private void resetMetrics() {
        startTime = System.currentTimeMillis();
        firstTokenTime = 0;
        totalTokens = 0;
        totalChars = 0;
    }
    
    private void appendMetrics(String content) {
        if (content == null || content.isEmpty()) return;
        
        if (showCharCount) {
            totalChars += content.length();
        }
        
        if (showTokenCount || showTokenCost) {
            // 简单估算：平均每个中文字符约2个token，英文单词约1.3个token
            double tokenEstimate = 0;
            for (char c : content.toCharArray()) {
                if (Character.UnicodeScript.of(c) == Character.UnicodeScript.HAN) {
                    tokenEstimate += 2.0;
                } else {
                    tokenEstimate += 1.3;
                }
            }
            totalTokens += Math.ceil(tokenEstimate);
        }
        
        if (showFirstTokenTime && firstTokenTime == 0) {
            firstTokenTime = System.currentTimeMillis();
        }
    }
    
    private void appendMessage(String message) {
        if (message == null || message.isEmpty()) {
            return;
        }
        
        Document doc = chatArea.getDocument();
        try {
            if (message.startsWith("AI: ") || message.startsWith("你: ")) {
                if (message.startsWith("你: ")) {
                    // 用户消息使用蓝色
                    doc.insertString(doc.getLength(), "\n", null);
                    chatArea.setForeground(new Color(0, 0, 255)); // 蓝色
                    // 加粗并放大显示"你: "
                    Font originalFont = chatArea.getFont();
                    Font boldLargeFont = originalFont.deriveFont(Font.BOLD, originalFont.getSize() + 2);
                    chatArea.setFont(boldLargeFont);
                    doc.insertString(doc.getLength(), "你: ", null);
                    chatArea.setFont(originalFont);
                    doc.insertString(doc.getLength(), message.substring(3) + "\n", null);
                    
                    // 如果是第一条消息，直接用作标题
                    ChatSession currentSession = sessionManager.getCurrentSession();
                    if (currentSession.getMessages().isEmpty()) {
                        String title = message.substring(3).trim(); // 去掉"你: "前缀
                        if (title.length() > 30) {
                            title = title.substring(0, 30) + "...";
                        }
                        currentSession.setTitle(title);
                        sessionManager.saveSession(currentSession);
                        updateSessionList();
                    }
                } else {
                    // AI消息使用绿色
                    doc.insertString(doc.getLength(), "\n", null);
                    chatArea.setForeground(new Color(0, 128, 0)); // 绿色
                    
                    if (message.equals("AI: ")) {
                        // 加粗并放大显示"AI (deepseek:latest): "
                        Font originalFont = chatArea.getFont();
                        Font boldLargeFont = originalFont.deriveFont(Font.BOLD, originalFont.getSize() + 2);
                        chatArea.setFont(boldLargeFont);
                        doc.insertString(doc.getLength(), "AI (" + selectedModel + "): ", null);
                        chatArea.setFont(originalFont);
                    } else {
                        // 加粗并放大显示"AI: "
                        Font originalFont = chatArea.getFont();
                        Font boldLargeFont = originalFont.deriveFont(Font.BOLD, originalFont.getSize() + 2);
                        chatArea.setFont(boldLargeFont);
                        doc.insertString(doc.getLength(), "AI: ", null);
                        chatArea.setFont(originalFont);
                        doc.insertString(doc.getLength(), message.substring(3), null);
                    }
                }
            } else {
                // 其他消息使用默认颜色
                chatArea.setForeground(new Color(0, 0, 0)); // 黑色
                doc.insertString(doc.getLength(), message, null);
            }
            
            // 保存消息到当前会话
            ChatSession currentSession = sessionManager.getCurrentSession();
            if (message.startsWith("AI: ") || message.startsWith("你: ")) {
                String role = message.startsWith("AI: ") ? "AI" : "用户";
                String content = message.substring(message.indexOf(": ") + 2);
                if (!content.isEmpty()) {
                    currentSession.addMessage(role, content);
                    sessionManager.saveSession(currentSession);
                }
            }
        } catch (BadLocationException e) {
            e.printStackTrace();
        }
    }

    private void appendToChat(String role, String message) {
        if (role.equals("AI")) {
            if (!isAIResponseStarted) {
                appendMessage("AI: ");
                isAIResponseStarted = true;
            }
            appendMessage(message);
        } else {
            isAIResponseStarted = false;
            appendMessage(role + ": " + message);
        }
    }
    
    private void sendMessage() {
        String message = inputField.getText().trim();
        if (message.isEmpty()) {
            return;
        }
        
        // 获取当前会话
        final ChatSession currentSession = sessionManager.getCurrentSession();
        if (currentSession == null) {
            appendToChat("系统", "错误：当前会话无效");
            return;
        }
            // 如果是第一条消息，自动更新会话标题
        if ("新会话".equals(currentSession.getTitle())) {
            currentSession.setTitle(message);
            sessionManager.saveSession(currentSession);
            updateSessionList();
        }
        // 禁用输入,避免重复发送
        inputField.setEnabled(false);
        sendButton.setEnabled(false);
        isAIResponseStarted = false;  // 重置状态
        messageBuffer.setLength(0);    // 清空消息缓冲区
        isFirstResponse.set(true);     // 重置首次响应标志
        
        // 清空输入框并显示用户消息
        inputField.setText("");
        appendToChat("你", message);
        
        // 使用线程池处理消息发送
        executorService.submit(() -> {
            try {
                // 再次验证当前会话
                if (sessionManager.getCurrentSession() != currentSession) {
                    SwingUtilities.invokeLater(() -> {
                        appendToChat("系统", "会话已切换，消息发送已取消");
                        inputField.setEnabled(true);
                        sendButton.setEnabled(true);
                        inputField.requestFocus();
                    });
                    return;
                }
                
                // 保存用户消息到当前会话
                currentSession.addMessage("用户", message);
                sessionManager.saveSession(currentSession);
                
                // 优化历史消息处理
                List<OllamaClient.Message> history = new ArrayList<>();
                int startIndex = Math.max(0, currentSession.getMessages().size() - maxContext + 1);
                List<ChatSession.ChatMessage> messages = currentSession.getMessages().subList(startIndex, currentSession.getMessages().size());
                
                // 批量转换历史消息
                for (ChatSession.ChatMessage msg : messages) {
                    history.add(new OllamaClient.Message(
                        msg.getRole().equals("用户") ? "user" : "assistant",
                        msg.getContent()
                    ));
                }
                
                // 取消之前的订阅
                if (currentSubscription != null && !currentSubscription.isDisposed()) {
                    currentSubscription.dispose();
                }
                
                final String finalMessage = message;
                StringBuilder responseBuilder = new StringBuilder();
                
                currentSubscription = ollamaClient.prompt()
                        .withUserMessage(finalMessage)
                        .withModel(selectedModel)
                        .withTemperature(temperature)
                        .withMaxContext(maxContext)
                        .withHistory(history)
                        .withSystem(spec -> {
                            spec.param("role", "assistant");
                            spec.param("name", "myllm");
                        })
                        .stream()
                        .subscribe(
                            response -> {
                                // 验证当前会话是否仍然有效
                                if (sessionManager.getCurrentSession() != currentSession) {
                                    if (currentSubscription != null && !currentSubscription.isDisposed()) {
                                        currentSubscription.dispose();
                                    }
                                    return;
                                }
                                
                                responseBuilder.append(response);
                                messageBuffer.append(response);
                                
                                // 当缓冲区达到一定大小时更新UI
                                if (messageBuffer.length() >= BUFFER_SIZE) {
                                    String bufferContent = messageBuffer.toString();
                                    messageBuffer.setLength(0);
                                    
                                    SwingUtilities.invokeLater(() -> {
                                        // 验证当前会话是否仍然有效
                                        if (sessionManager.getCurrentSession() != currentSession) {
                                            return;
                                        }
                                        
                                        if (isFirstResponse.get()) {
                                            appendToChat("AI", bufferContent);
                                            isFirstResponse.set(false);
                                        } else {
                                            appendMessage(bufferContent);
                                        }
                                    });
                                }
                            },
                            error -> {
                                SwingUtilities.invokeLater(() -> {
                                    // 验证当前会话是否仍然有效
                                    if (sessionManager.getCurrentSession() != currentSession) {
                                        return;
                                    }
                                    
                                    appendToChat("系统", "错误: " + error.getMessage());
                                    inputField.setEnabled(true);
                                    sendButton.setEnabled(true);
                                    inputField.requestFocus();
                                });
                            },
                            () -> {
                                // 验证当前会话是否仍然有效
                                if (sessionManager.getCurrentSession() != currentSession) {
                                    return;
                                }
                                
                                // 处理剩余的缓冲区内容
                                if (messageBuffer.length() > 0) {
                                    String remainingContent = messageBuffer.toString();
                                    messageBuffer.setLength(0);
                                    
                                    SwingUtilities.invokeLater(() -> {
                                        if (isFirstResponse.get()) {
                                            appendToChat("AI", remainingContent);
                                            isFirstResponse.set(false);
                                        } else {
                                            appendMessage(remainingContent);
                                        }
                                    });
                                }
                                
                                // 在完成时保存完整的响应
                                SwingUtilities.invokeLater(() -> {
                                    String fullResponse = responseBuilder.toString();
                                    if (!fullResponse.isEmpty()) {
                                        currentSession.addMessage("AI", fullResponse);
                                        currentSession.setLastResponseTime(System.currentTimeMillis());
                                        sessionManager.saveSession(currentSession);
                                        updateSessionList(); // 更新会话列表
                                    }
                                    appendMetrics("");
                                    inputField.setEnabled(true);
                                    sendButton.setEnabled(true);
                                    inputField.requestFocus();
                                });
                            }
                        );
            } catch (Exception e) {
                SwingUtilities.invokeLater(() -> {
                    // 验证当前会话是否仍然有效
                    if (sessionManager.getCurrentSession() != currentSession) {
                        return;
                    }
                    
                    appendToChat("系统", "发送消息时发生错误: " + e.getMessage());
                    inputField.setEnabled(true);
                    sendButton.setEnabled(true);
                    inputField.requestFocus();
                });
            }
        });
    }

    private String handleLongText(String text) {
        try {
            File tempFile = File.createTempFile("myllm_long_text_", ".txt");
            try (PrintWriter writer = new PrintWriter(tempFile)) {
                writer.write(text);
            }
            return String.format("请分析以下文件内容：%s", tempFile.getAbsolutePath());
        } catch (IOException e) {
            e.printStackTrace();
            return text;
        }
    }

    private String processContent(String content) {
        if (content == null || content.isEmpty()) return content;
        
        // 拼写检查
        if (spellCheck) {
            content = performSpellCheck(content);
        }
        
        // Markdown 渲染
        if (markdownRendering) {
            content = renderMarkdown(content);
            
            // LaTeX 渲染
            if (latexRendering) {
                content = renderLatex(content);
            }
            
            // Mermaid 渲染
            if (mermaidRendering) {
                content = renderMermaid(content);
            }
        }
        
        return content;
    }

    private String performSpellCheck(String text) {
        // 简单的拼写检查实现
        // TODO: 接入专业的拼写检查库
        return text;
    }

    private String renderMarkdown(String text) {
        // 简单的 Markdown 渲染
        // TODO: 使用专业的 Markdown 渲染库
        return text;
    }

    private String renderLatex(String text) {
        // 简单的 LaTeX 渲染
        // TODO: 使用专业的 LaTeX 渲染库
        return text;
    }

    private String renderMermaid(String text) {
        // 简单的 Mermaid 渲染
        // TODO: 使用专业的 Mermaid 渲染库
        return text;
    }

    private void generateTitle(String message) {
        try {
            ollamaClient.prompt()
                    .withUserMessage(message)
                    .withModel("deepseek:latest")
                    .withTemperature(0.7)
                    .withMaxContext(4096)
                    .withSystem(spec -> {
                        spec.param("role", "assistant");
                        spec.param("name", "myllm");
                    })
                    .stream()
                    .subscribe(
                        title -> {
                            ChatSession session = sessionManager.getCurrentSession();
                            session.setTitle(title);
                            sessionManager.saveSession(session);
                            updateSessionList();
                        },
                        error -> {
                            appendToChat("系统", "生成标题时出错: " + error.getMessage());
                        },
                        () -> {
                            System.out.println("标题生成完成");
                        }
                    );
        } catch (IOException e) {
            appendToChat("系统", "生成标题时发生错误: " + e.getMessage());
        }
    }

    private void createNewSession() {
        ChatSession newSession = sessionManager.createNewSession();
        updateSessionList();
        sessionList.setSelectedValue(newSession, true);
        switchSession(newSession);
    }

    private void switchSession(ChatSession session) {
        if (session == null) {
            System.err.println("尝试切换到空会话");
            return;
        }

        // 如果是同一个会话，不需要切换
        if (session.getId().equals(currentSessionId)) {
            System.out.println("已经是当前会话，无需切换");
            return;
        }

        System.out.println("切换到会话: " + session.getTitle() + " (ID: " + session.getId() + ")");
        
        // 保存当前会话的响应状态
        ChatSession currentSession = sessionManager.getCurrentSession();
        if (currentSession != null && currentSubscription != null && !currentSubscription.isDisposed()) {
            currentSession.setLastResponseTime(System.currentTimeMillis());
            sessionManager.saveSession(currentSession);
        }
        
        // 取消当前正在进行的对话
        if (currentSubscription != null && !currentSubscription.isDisposed()) {
            currentSubscription.dispose();
            currentSubscription = null;
        }
        
        // 清空消息缓冲区
        messageBuffer.setLength(0);
        isFirstResponse.set(true);
        isAIResponseStarted = false;
        
        // 更新当前会话
        sessionManager.setCurrentSession(session);
        currentSessionId = session.getId();
        
        // 在EDT线程中更新UI
        SwingUtilities.invokeLater(() -> {
            try {
                // 清空聊天区域
                chatArea.setText("");
                
                // 重新显示历史消息
                List<ChatSession.ChatMessage> messages = session.getMessages();
                for (int i = 0; i < messages.size(); i++) {
                    ChatSession.ChatMessage message = messages.get(i);
                    String role = message.getRole();
                    String content = message.getContent();
                    
                    if (role.equals("AI")) {
                        if (i == 0 || !messages.get(i-1).getRole().equals("AI")) {
                            appendMessage("AI: ");
                        }
                        appendMessage(content);
                    } else {
                        appendMessage(role + ": " + content);
                    }
                }
                
                // 确保输入框可用
                inputField.setEnabled(true);
                sendButton.setEnabled(true);
                inputField.requestFocus();
                
                // 更新会话列表
                updateSessionList();
                
                // 滚动到底部
                chatArea.setCaretPosition(chatArea.getDocument().getLength());
                
                // 强制重绘
                chatArea.repaint();
                
            } catch (Exception e) {
                System.err.println("更新UI时发生错误: " + e.getMessage());
                e.printStackTrace();
            }
        });
    }

    private void updateSessionList() {
        SwingUtilities.invokeLater(() -> {
            sessionListModel.clear();
            List<ChatSession> allSessions = sessionManager.getAllSessions();
            for (ChatSession session : allSessions) {
                sessionListModel.addElement(session);
            }
            
            // 确保当前会话被选中
            ChatSession current = sessionManager.getCurrentSession();
            if (current != null) {
                sessionList.setSelectedValue(current, true);
            }
        });
    }

    private void loadCurrentSession() {
        ChatSession current = sessionManager.getCurrentSession();
        sessionList.setSelectedValue(current, true);
        switchSession(current);
    }
    
    private void showSettings() {
        JDialog dialog = new JDialog(this, "设置", true);
        dialog.setSize(500, 500);
        dialog.setLocationRelativeTo(this);

        JTabbedPane tabbedPane = new JTabbedPane();
        
        // API设置面板
        JPanel apiPanel = new JPanel(new GridBagLayout());
        GridBagConstraints gbc = new GridBagConstraints();
        gbc.fill = GridBagConstraints.HORIZONTAL;
        gbc.insets = new Insets(5, 5, 5, 5);
        
        JTextField apiEndpointField = new JTextField(apiEndpoint);
        JTextField modelField = new JTextField(selectedModel);
        JSpinner maxContextSpinner = new JSpinner(new SpinnerNumberModel(maxContext, 1, 32768, 1));
        JSpinner temperatureSpinner = new JSpinner(new SpinnerNumberModel(temperature, 0.0, 2.0, 0.1));
        
        gbc.gridx = 0; gbc.gridy = 0;
        apiPanel.add(new JLabel("API地址:"), gbc);
        gbc.gridx = 1;
        apiPanel.add(apiEndpointField, gbc);
        
        gbc.gridx = 0; gbc.gridy = 1;
        apiPanel.add(new JLabel("模型:"), gbc);
        gbc.gridx = 1;
        apiPanel.add(modelField, gbc);
        
        gbc.gridx = 0; gbc.gridy = 2;
        apiPanel.add(new JLabel("最大上下文:"), gbc);
        gbc.gridx = 1;
        apiPanel.add(maxContextSpinner, gbc);
        
        gbc.gridx = 0; gbc.gridy = 3;
        apiPanel.add(new JLabel("温度:"), gbc);
        gbc.gridx = 1;
        apiPanel.add(temperatureSpinner, gbc);
        
        // 对话设置面板
        JPanel conversationPanel = new JPanel(new GridBagLayout());
        gbc = new GridBagConstraints();
        gbc.fill = GridBagConstraints.HORIZONTAL;
        gbc.insets = new Insets(5, 5, 5, 5);
        
        JTextArea systemPromptArea = new JTextArea(systemPrompt, 5, 30);
        systemPromptArea.setLineWrap(true);
        systemPromptArea.setWrapStyleWord(true);
        
        JCheckBox autoTitleBox = new JCheckBox("自动生成标题", autoGenerateTitle);
        JCheckBox spellCheckBox = new JCheckBox("拼写检查", spellCheck);
        JCheckBox markdownBox = new JCheckBox("Markdown渲染", markdownRendering);
        JCheckBox latexBox = new JCheckBox("LaTeX渲染", latexRendering);
        JCheckBox mermaidBox = new JCheckBox("Mermaid渲染", mermaidRendering);
        JCheckBox injectDataBox = new JCheckBox("注入默认数据", injectDefaultData);
        JCheckBox artifactsBox = new JCheckBox("自动生成工件", autoGenerateArtifacts);
        JCheckBox longTextBox = new JCheckBox("长文本转文件", longTextAsFile);
        
        gbc.gridx = 0; gbc.gridy = 0;
        gbc.gridwidth = 2;
        conversationPanel.add(new JLabel("系统提示:"), gbc);
        
        gbc.gridy = 1;
        conversationPanel.add(new JScrollPane(systemPromptArea), gbc);
        
        gbc.gridy = 2;
        conversationPanel.add(autoTitleBox, gbc);
        
        gbc.gridy = 3;
        conversationPanel.add(spellCheckBox, gbc);
        
        gbc.gridy = 4;
        conversationPanel.add(markdownBox, gbc);
        
        gbc.gridy = 5;
        conversationPanel.add(latexBox, gbc);
        
        gbc.gridy = 6;
        conversationPanel.add(mermaidBox, gbc);
        
        gbc.gridy = 7;
        conversationPanel.add(injectDataBox, gbc);
        
        gbc.gridy = 8;
        conversationPanel.add(artifactsBox, gbc);
        
        gbc.gridy = 9;
        conversationPanel.add(longTextBox, gbc);
        
        tabbedPane.addTab("API设置", apiPanel);
        tabbedPane.addTab("对话设置", conversationPanel);
        
        JPanel buttonPanel = new JPanel(new FlowLayout(FlowLayout.RIGHT));
        JButton saveButton = new JButton("保存");
        JButton cancelButton = new JButton("取消");
        
        saveButton.addActionListener(e -> {
            prefs.put("apiEndpoint", apiEndpoint = apiEndpointField.getText());
            prefs.put("selectedModel", selectedModel = modelField.getText());
            prefs.putInt("maxContext", maxContext = (Integer) maxContextSpinner.getValue());
            prefs.putDouble("temperature", temperature = (Double) temperatureSpinner.getValue());
            prefs.put("systemPrompt", systemPrompt = systemPromptArea.getText());
            prefs.putBoolean("autoGenerateTitle", autoGenerateTitle = autoTitleBox.isSelected());
            prefs.putBoolean("spellCheck", spellCheck = spellCheckBox.isSelected());
            prefs.putBoolean("markdownRendering", markdownRendering = markdownBox.isSelected());
            prefs.putBoolean("latexRendering", latexRendering = latexBox.isSelected());
            prefs.putBoolean("mermaidRendering", mermaidRendering = mermaidBox.isSelected());
            prefs.putBoolean("injectDefaultData", injectDefaultData = injectDataBox.isSelected());
            prefs.putBoolean("autoGenerateArtifacts", autoGenerateArtifacts = artifactsBox.isSelected());
            prefs.putBoolean("longTextAsFile", longTextAsFile = longTextBox.isSelected());
            
            ollamaClient = new OllamaClient(apiEndpoint);
            dialog.dispose();
        });
        
        cancelButton.addActionListener(e -> dialog.dispose());
        
        buttonPanel.add(saveButton);
        buttonPanel.add(cancelButton);
        
        dialog.setLayout(new BorderLayout());
        dialog.add(tabbedPane, BorderLayout.CENTER);
        dialog.add(buttonPanel, BorderLayout.SOUTH);
        dialog.setVisible(true);
    }

    private void showPerformanceMonitor() {
        JDialog dialog = new JDialog(this, "性能监控", false);  // 使用非模态对话框
        dialog.setSize(600, 400);
        dialog.setLocationRelativeTo(this);
        
        JTextArea monitorArea = new JTextArea();
        monitorArea.setEditable(false);
        JScrollPane scrollPane = new JScrollPane(monitorArea);
        
        JPanel controlPanel = new JPanel(new FlowLayout(FlowLayout.LEFT));
        JCheckBox autoRefresh = new JCheckBox("自动刷新", true);
        JButton refreshButton = new JButton("刷新");
        controlPanel.add(autoRefresh);
        controlPanel.add(refreshButton);
        
        Timer timer = new Timer(1000, e -> updatePerformanceInfo(monitorArea));
        timer.start();
        
        autoRefresh.addActionListener(e -> {
            if (autoRefresh.isSelected()) {
                timer.start();
            } else {
                timer.stop();
            }
        });
        
        refreshButton.addActionListener(e -> updatePerformanceInfo(monitorArea));
        
        dialog.addWindowListener(new WindowAdapter() {
            @Override
            public void windowClosing(WindowEvent e) {
                timer.stop();
            }
        });
        
        dialog.setLayout(new BorderLayout());
        dialog.add(controlPanel, BorderLayout.NORTH);
        dialog.add(scrollPane, BorderLayout.CENTER);
        dialog.setVisible(true);
    }

    private void updatePerformanceInfo(JTextArea area) {
        StringBuilder info = new StringBuilder();
        Runtime runtime = Runtime.getRuntime();
        
        // 内存使用情况
        long totalMemory = runtime.totalMemory();
        long freeMemory = runtime.freeMemory();
        long usedMemory = totalMemory - freeMemory;
        long maxMemory = runtime.maxMemory();
        
        info.append("=== 内存使用情况 ===\n");
        info.append(String.format("最大可用内存: %.2f MB\n", maxMemory / 1024.0 / 1024.0));
        info.append(String.format("已分配内存: %.2f MB\n", totalMemory / 1024.0 / 1024.0));
        info.append(String.format("已使用内存: %.2f MB\n", usedMemory / 1024.0 / 1024.0));
        info.append(String.format("空闲内存: %.2f MB\n", freeMemory / 1024.0 / 1024.0));
        info.append("\n");
        
        // 线程信息
        info.append("=== 线程信息 ===\n");
        ThreadGroup root = Thread.currentThread().getThreadGroup().getParent();
        while (root.getParent() != null) {
            root = root.getParent();
        }
        
        Thread[] threads = new Thread[root.activeCount()];
        root.enumerate(threads);
        
        int running = 0, waiting = 0, blocked = 0;
        for (Thread t : threads) {
            if (t != null) {
                switch (t.getState()) {
                    case RUNNABLE: running++; break;
                    case WAITING:
                    case TIMED_WAITING: waiting++; break;
                    case BLOCKED: blocked++; break;
                }
            }
        }
        
        info.append(String.format("运行中线程: %d\n", running));
        info.append(String.format("等待中线程: %d\n", waiting));
        info.append(String.format("阻塞线程: %d\n", blocked));
        info.append("\n");
        
        // 会话信息
        info.append("=== 会话信息 ===\n");
        info.append(String.format("会话总数: %d\n", sessionManager.getAllSessions().size()));
        ChatSession current = sessionManager.getCurrentSession();
        if (current != null) {
            info.append(String.format("当前会话消息数: %d\n", current.getMessages().size()));
            info.append(String.format("当前会话标题: %s\n", current.getTitle()));
        }
        info.append("\n");
        
        // UI响应时间
        info.append("=== UI响应时间 ===\n");
        long uiResponseTime = System.nanoTime();
        SwingUtilities.invokeLater(() -> {
            long responseDelay = (System.nanoTime() - uiResponseTime) / 1000000;
            info.append(String.format("UI线程响应延迟: %d ms\n", responseDelay));
            area.setText(info.toString());
        });
    }

    // 优化消息处理
    private void optimizeMessageProcessing() {
        // 清理过期的会话
        List<ChatSession> sessions = sessionManager.getAllSessions();
        long now = System.currentTimeMillis();
        for (ChatSession session : sessions) {
            if (session != sessionManager.getCurrentSession() && 
                (now - session.getLastResponseTime()) > 30 * 60 * 1000) { // 30分钟无活动
                sessionManager.deleteSession(session);
            }
        }
        
        // 定期清理内存
        if (Runtime.getRuntime().freeMemory() < Runtime.getRuntime().totalMemory() * 0.2) { // 空闲内存少于20%
            System.gc();
        }
    }

    private class SessionListCellRenderer extends DefaultListCellRenderer {
        @Override
        public Component getListCellRendererComponent(JList<?> list, Object value, int index,
                                                    boolean isSelected, boolean cellHasFocus) {
            super.getListCellRendererComponent(list, value, index, isSelected, cellHasFocus);
            if (value instanceof ChatSession) {
                ChatSession session = (ChatSession) value;
                String title = session.getTitle();
                if (title == null || title.isEmpty()) {
                    title = "新会话 " + session.getId().substring(0, 8);
                }
                if (session.isPinned()) {
                    title = "📌 " + title;
                }
                
                // 添加响应状态指示
                if (session == sessionManager.getCurrentSession() && 
                    currentSubscription != null && 
                    !currentSubscription.isDisposed()) {
                    title = "💬 " + title;
                }
                
                setText(title);
                
                // 添加鼠标监听器支持双击编辑标题
                if (list.getMouseListeners().length == 0) {
                    list.addMouseListener(new MouseAdapter() {
                        public void mouseClicked(MouseEvent e) {
                            if (e.getClickCount() == 2) {
                                int index = list.locationToIndex(e.getPoint());
                                ChatSession session = sessionListModel.getElementAt(index);
                                String newTitle = (String) JOptionPane.showInputDialog(
                                        MyllmClient.this,
                                    "请输入新的标题：",
                                    "编辑标题",
                                    JOptionPane.PLAIN_MESSAGE,
                                    null,
                                    null,
                                    session.getTitle()
                                );
                                if (newTitle != null && !newTitle.trim().isEmpty()) {
                                    session.setTitle(newTitle.trim());
                                    sessionManager.saveSession(session);
                                    updateSessionList();
                                }
                            }
                        }
                    });
                }
            }
            return this;
        }
    }
    
    public static void main(String[] args) {
            try {
                UIManager.setLookAndFeel(UIManager.getSystemLookAndFeelClassName());
            } catch (Exception e) {
                e.printStackTrace();
            }
        
        SwingUtilities.invokeLater(() -> {
            new MyllmClient().setVisible(true);
        });
    }

    @Override
    public void dispose() {
        executorService.shutdown();
        try {
            if (!executorService.awaitTermination(5, TimeUnit.SECONDS)) {
                executorService.shutdownNow();
            }
        } catch (InterruptedException e) {
            executorService.shutdownNow();
            Thread.currentThread().interrupt();
        }
        super.dispose();
    }
} 